import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Activation
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import LSTM, TimeDistributed
from tensorflow.keras import utils
# 时间步：一个字母一个时间步

sample = "hihello"

# 模型只能计算数字，字母转化为数字
char_set = list(set(sample))  # id -> char ['i', 'l', 'e', 'o', 'h']
char_dic = {w: i for i, w in enumerate(char_set)}

# x: hihell
# y: ihello
# 输入：hihell
# 输出：ihello
x_str = sample[:-1]
y_str = sample[1:]

#定义rnn模型参数
data_dim = len(char_set)  # 输入数据维度:有多少个字母，多少个字符类别=输入数据维度
timesteps = len(y_str)    # 序列长度:时间步， y标签长度， 每个字母一个时间步
num_classes = len(char_set) # 字符类别个数=y的种类

print(x_str, y_str)

x = [char_dic[c] for c in x_str]  # 字符转化为数字形式  [6]
y = [char_dic[c] for c in y_str]  # char to index

# 输入进行独热编码：One-hot encoding
x = utils.to_categorical(x, num_classes=num_classes) # [6, 5]
# 将x变成rnn需要的输入的shape:(样本个数,序列长度,输入数据纬度) reshape X to be [samples, time steps, features]
x = np.reshape(x, (-1, len(x), data_dim))
print(x.shape) # [ 1, 6, 5]

# 输出y需要独热编码(多分类)：One-hot encoding
y = utils.to_categorical(y, num_classes=num_classes)
# 将y变成rnn需要的输出的shape: (样本个数,序列长度,输入数据纬度)
y = np.reshape(y, (-1, len(y), data_dim))
print(y.shape)
#建立模型
model = Sequential()
#构建LSTM网络：num_classes:输出字符种类(隐藏层单元个数); input_shape输入维度(序列长度,输入数据维度)；return_sequences：返回各序列=true
model.add(LSTM(num_classes, input_shape=(timesteps, data_dim), return_sequences=True))
#TimeDistributed:按时间分布构筑输出网络
model.add(TimeDistributed(Dense(num_classes))) # 6，5  5*6,1
model.add(Activation('softmax'))
model.summary()

model.compile(loss='categorical_crossentropy', #多分类:需要独热
              optimizer='rmsprop', metrics=['accuracy'])
model.fit(x, y, epochs=800)

predictions = model.predict(x, verbose=0)# [1, 6, 5]
print('---\n', predictions)
for i, prediction in enumerate(predictions): # prediction [6， 5]
    x_index = np.argmax(x[i], axis=1)   #输入x
    x_str = [char_set[j] for j in x_index] #通过id找到字符
    print(x_index, ''.join(x_str))

    index = np.argmax(prediction, axis=1)   #输出预测值prediction
    result = [char_set[j] for j in index]  #通过id找到字符
    print(index, ''.join(result))
